# classes
import torch
from torch.utils.data import Dataset
from typing import Tuple
import pandas as pd
import numpy as np
import logging
import os

CONFIG = {
    "data_path": r"C:/智能螺栓/螺栓实时数据监测.csv",
    "batch_size": 256,
    "epochs": 3000,
    "lr": 1e-3,
    "weight_decay": 1e-2,
    "save_dir": "C:/智能螺栓/bolt_model_v7",
    "device": "cuda" if torch.cuda.is_available() else "cpu",
    "plot_interval": 50 
}

def setup_logger():
    logger = logging.getLogger("BoltTrainer")
    logger.setLevel(logging.INFO)    
    log_dir = CONFIG["save_dir"]
    os.makedirs(log_dir, exist_ok=True)    
    file_handler = logging.FileHandler(os.path.join(log_dir, "train.log"))
    stream_handler = logging.StreamHandler()   
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    file_handler.setFormatter(formatter)
    stream_handler.setFormatter(formatter)  
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)
    return logger
log = setup_logger()

class BoltDataset(Dataset):
    '''数据集'''

    def __init__(self):
        self.df = pd.read_csv(CONFIG["data_path"])
        self._preprocess()
        self._init_normalizers()

    def _preprocess(self):
        q1_f = self.df["预紧力F(N)"].quantile(0.25)
        q3_f = self.df["预紧力F(N)"].quantile(0.75)
        iqr_f = q3_f - q1_f
        self.df = self.df[(self.df["预紧力F(N)"] >= q1_f - 1.5 * iqr_f) & 
                         (self.df["预紧力F(N)"] <= q3_f + 1.5 * iqr_f)]        
        q1_v = self.df["实测电压V(V)"].quantile(0.25)
        q3_v = self.df["实测电压V(V)"].quantile(0.75)
        iqr_v = q3_v - q1_v
        self.df = self.df[(self.df["实测电压V(V)"] >= q1_v - 1.5 * iqr_v) & 
                         (self.df["实测电压V(V)"] <= q3_v + 1.5 * iqr_v)]        
        self.df.reset_index(drop=True, inplace=True)

    def _init_normalizers(self):
        self.V_mean = self.df["实测电压V(V)"].mean()
        self.V_std = self.df["实测电压V(V)"].std() or 1.0
        self.F_mean = self.df["预紧力F(N)"].mean()
        self.F_std = self.df["预紧力F(N)"].std() or 1.0

    def normalize_V(self, V: np.ndarray) -> np.ndarray:
        return (V - self.V_mean) / self.V_std
    
    def denormalize_F(self, F_norm: np.ndarray) -> np.ndarray:
        return F_norm * self.F_std + self.F_mean
    
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, ...]:
        row = self.df.iloc[idx]
        V_norm = self.normalize_V(np.array([row["实测电压V(V)"]]))[0]
        F_norm = (row["预紧力F(N)"] - self.F_mean) / self.F_std
        
        A = row["横截面积A(m²)"]
        E = row["弹性模量E(Pa)"]
        g31 = row["压电常数g31(V·m/N)"]
        t = row["压电片厚度t(m)"]
        V_raw = row["实测电压V(V)"]       

        return (
            torch.tensor([V_norm], dtype=torch.float32),
            torch.tensor([F_norm], dtype=torch.float32),
            torch.tensor([A], dtype=torch.float32),
            torch.tensor([E], dtype=torch.float32),
            torch.tensor([g31], dtype=torch.float32),
            torch.tensor([t], dtype=torch.float32),
            torch.tensor([V_raw], dtype=torch.float32)
        )